#pragma once
#include "kernel_const.h"
#include "kernel_utils.h"
#include "kernel_operator.h"


/**
 * @brief 用于 GM内ND行的AB->L1内Nz的AB
GM中的数据不要求对齐，但是最大为(L1_M0, L1_N0, L1_K0)
L1目的地址得到的数据应该GM内有效数据按Mmad最小基块的大小（baseblock的M0N0K0）向上对齐后的数据。
读入的除有效数据以外其他为L1内脏数据。
 */

[aicore] inline __attribute__((always_inline)) void Gm2L1_Nd2Nz(
    AscendC::LocalTensor<half> l1_dstLocalTensor,
    AscendC::GlobalTensor<half> gm_srcGlobalTensor,
    uint32_t gm_mActual,            // 实际需要的数据大小。最大为(L1_M0, L1_N0, L1_K0)
    uint32_t gm_nActual,
    uint32_t gm_srcStride,
    uint32_t l1_mActualRound,       // 实际需要的数据大小对基块维度(BASEBLOCK_M0, BASEBLOCK_N0=向上取整。（这里传入的数应该是按BASEBLOCK_M0或N0对齐了的，对于half，BASEBLOCK_M0N0K0 3个数相等）
    uint32_t l1_nActualRound        // 这个参数没有用，N方向是自动向32B（16个half元素）去对齐的，这里是因为Nd2NzParams的参数dstNzNStride的单位为32B（外面传入的应该是对BASEBLOCK_K0对齐过的数）
){
    if( gm_srcStride < STRIDE_LIMIT ){
        AscendC::DataCopy(
            l1_dstLocalTensor, 
            gm_srcGlobalTensor, 
            AscendC::Nd2NzParams(
                1, 
                gm_mActual, 
                gm_nActual, 
                0, 
                gm_srcStride, 
                l1_mActualRound, 
                1, 
                1 // 0？
            )
        );
    } else {
        for(uint64_t i = 0; i < gm_mActual; i++ ){
            AscendC::DataCopy(
                l1_dstLocalTensor[i*DATABLOCK_SIZE], //datablock是z一行或n一列的大小
                gm_srcGlobalTensor[i*gm_srcStride],
                AscendC::Nd2NzParams(
                    1, 
                    1, 
                    gm_nActual, 
                    0, 
                    gm_srcStride, 
                    l1_mActualRound, 
                    0, 
                    1 // 0？
                )  
            );
        }
    }
}

/**
 * @brief 用于 L1内Nz的A->L0A内Zz的A
这里传入的一定是按Mmad基块对齐过的数据，在GM->L1的时候，到L1就已对齐，可能有L1内的无效数据，这里会搬入L0。
计算完后搬出时只搬运需要的即可。
搬运前需要先计算L1的(M0, N0, K0)块需要划分多少次搬入L0计算，只需要计算搬入L0的K维度循环几次即可，MN维度与L1相同
 */
[aicore] inline __attribute__((always_inline)) void L12L0_Nz2Zz(
    AscendC::LocalTensor<half> l0_dstLocalTensor, 
    AscendC::LocalTensor<half> l1_srcLocalTensor, 
    uint32_t l1_mActualRound, 
    uint32_t l1_kActualRound_loop, 
    uint32_t l1_srcStride,  // 这里等于l1_mActualRound
    uint32_t l0_baseblockM, //Mmad最小基块参数
    uint32_t l0_baseblockK
){
    for(uint64_t i = 0; i < l1_mActualRound/l0_baseblockM; i++ ){
        AscendC::LoadData(
            l0_dstLocalTensor[i * l0_baseblockM * l1_kActualRound_loop], 
            l1_srcLocalTensor[i * l0_baseblockM * l0_baseblockK], //即每次跨一个基块大小
            AscendC::LoadData2DParams(
                0, // 注意这里的src，已经取过偏移了
                l1_kActualRound_loop / l0_baseblockK, 
                l1_mActualRound / l0_baseblockM, 
                0, 
                0, 
                false, 
                0
            )
        );
    }

}


[aicore] inline __attribute__((always_inline)) void L12L0_Nz2Zn(
    AscendC::LocalTensor<half> l0_dstLocalTensor, 
    AscendC::LocalTensor<half> l1_srcLocalTensor, 
    uint32_t l1_kActualRound_loop, 
    uint32_t l1_nActualRound, 
    uint32_t l1_srcStride,  // 注意这里用得上！传入l1_kActualRoundBaseblock
    uint32_t l0_baseblockK, //Mmad最小基块参数
    uint32_t l0_baseblockN
){
    for(uint64_t i = 0; i < l1_kActualRound_loop/l0_baseblockK; i++ ){
        AscendC::LoadData(
            l0_dstLocalTensor[i * l0_baseblockK * l1_nActualRound], 
            l1_srcLocalTensor[i * l0_baseblockK * l0_baseblockN], 
            AscendC::LoadData2DParams(
                0, 
                l1_nActualRound / l0_baseblockN, 
                l1_srcStride / l0_baseblockK, 
                0, 
                0, 
                true, 
                0
            )
        );
    }
}

[aicore] inline __attribute__((always_inline)) void L0C2Gm_Nz2Nd(
    AscendC::GlobalTensor<half> gm_dstGlobalTensor, 
    AscendC::LocalTensor<float> l0c_srcLocalTensor, 
    uint32_t l0c_mActual,    //实际需要搬出的数据大小，对应GM搬入时的实际数据
    uint32_t l0c_nActual,
    uint32_t l1_mActualRound,   // GM搬入后在L1按基块向上取整了才计算
    uint32_t l1_nActualRound,
    uint32_t gm_dstStride   // 搬入到GM的C矩阵后每行之间的距离
){
    AscendC::DataCopy(
        gm_dstGlobalTensor, 
        l0c_srcLocalTensor, 
        AscendC::DataCopyCO12DstParams(
            l0c_nActual, 
            l0c_mActual, 
            gm_dstStride, 
            l1_mActualRound, 
            QuantMode_t::F322F16, 
            0, 
            false,
            true
        )
    );
//     AscendC::FixpipeParamsV220 params;
//     params.nSize = l0c_nActual;
//     params.mSize = l0c_mActual;
//     params.srcStride = 32 * l1_mActualRound / (16 * sizeof(half));
//     params.dstStride = gm_dstStride;
//     params.quantPre = QuantMode_t::F322F16;
//     params.ndNum = 1;
//     params.reluEn = false;
//     params.isChannelSplit = false;
//     AscendC::Fixpipe<half, float, AscendC::CFG_ROW_MAJOR>(
//         gm_dstGlobalTensor, 
//         l0c_srcLocalTensor, 
//         params
//     );
}

[aicore] inline __attribute__((always_inline)) void Gm2Ub(
    AscendC::LocalTensor<half> ub_dstLocalTensor, 
    AscendC::GlobalTensor<half> gm_srcGlobalTensor, 
    uint32_t gm_mActual, 
    uint32_t gm_nActual, 
    uint32_t gm_srcStride 
    // 搬入UB后各行有效数据首尾连续
){
    AscendC::DataCopyPad<half>(
        ub_dstLocalTensor, 
        gm_srcGlobalTensor, 
        AscendC::DataCopyExtParams(
            gm_mActual, 
            gm_nActual * sizeof(half), 
            // TODO这个参数传参都不能超过uint32_t表示范围
            (gm_srcStride - gm_nActual) * sizeof(half), 
            0,
            0
        ),
        AscendC::DataCopyPadExtParams<half>(
            false, 
            0, 
            0, 
            0
        )
    );
}

[aicore] inline __attribute__((always_inline)) void Ub2Gm(
    AscendC::GlobalTensor<half> gm_dstGlobalTensor, 
    AscendC::LocalTensor<half> ub_srcLocalTensor, 
    uint32_t ub_mActual, 
    uint32_t ub_nActual, 
    // UB内有效数据首尾连续
    uint32_t gm_dstStride
){
    AscendC::DataCopyPad<half>(
        gm_dstGlobalTensor, 
        ub_srcLocalTensor, 
        AscendC::DataCopyExtParams(
            ub_mActual, 
            ub_nActual * sizeof(half), 
            0, 
            // TODO这个参数传参都不能超过uint32_t表示范围
            (gm_dstStride - ub_nActual) * sizeof(half),
            0
        )
    );
}